import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torch.optim as optim
import math
from torchvision.transforms.functional import InterpolationMode
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


batch_size = 64
batch_size_ = 16
num_workers = 2
learning_rate=0.001
epochs = 50  # 50轮
# 判断是否有GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 图像预处理
mean = [0.5070751592371323, 0.48654887331495095, 0.4409178433670343]
std = [0.2673342858792401, 0.2564384629170883, 0.27615047132568404]

transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])
transform_test = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize(mean, std)])

transform_train_224 = transforms.Compose([
    transforms.Resize(256, interpolation=InterpolationMode.BICUBIC),
    transforms.RandomCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])
transform_test_224 = transforms.Compose(
    [transforms.Resize(256, interpolation=InterpolationMode.BICUBIC),
     transforms.CenterCrop(224),
     transforms.ToTensor(),
     transforms.Normalize(mean, std)])

# CIFAR-100 数据集下载
train_dataset = torchvision.datasets.CIFAR100(root='data/',
                                             train=True,
                                             transform=transform_train,
                                             download=False)

test_dataset = torchvision.datasets.CIFAR100(root='data/',
                                            train=False,
                                            transform=transform_test)

# 数据载入
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           num_workers=num_workers,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          num_workers=num_workers,
                                          shuffle=False)


# CIFAR-100 数据集下载
train_dataset_224 = torchvision.datasets.CIFAR100(root='data/',
                                             train=True,
                                             transform=transform_train_224,
                                             download=False)

test_dataset_224 = torchvision.datasets.CIFAR100(root='data/',
                                            train=False,
                                            transform=transform_test_224)

# 数据载入
train_loader_224 = torch.utils.data.DataLoader(dataset=train_dataset_224,
                                           batch_size=batch_size_,
                                           num_workers=num_workers,
                                           shuffle=True)

test_loader_224 = torch.utils.data.DataLoader(dataset=test_dataset_224,
                                          batch_size=batch_size_,
                                          num_workers=num_workers,
                                          shuffle=False)
